import cv2
import numpy as np
import glob
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 1. 定义棋盘格参数（内部角点行列数）
rows = 6  # 棋盘格行数（内部角点数量）
cols = 9  # 棋盘格列数（内部角点数量）
square_size = 20  # 每个方格的物理尺寸（单位：任意，如毫米）

# 2. 生成棋盘格3D坐标 (Z=0)
objp = np.zeros((rows * cols, 3), np.float32)
objp[:, :2] = np.mgrid[0:cols, 0:rows].T.reshape(-1, 2) * square_size

# 3. 存储3D和2D点
obj_points = []  # 世界坐标系中的3D点
img_points = []  # 图像坐标系中的2D点

# 4. 读取所有标定图片
images = glob.glob('./chess1.jpg')  # 替换为你的图片路径

for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 5. 检测棋盘格角点
    ret, corners = cv2.findChessboardCorners(gray, (cols, rows), None)

    if ret:
        # 6. 亚像素精确化角点
        corners_subpix = cv2.cornerSubPix(
            gray, corners, (11, 11), (-1, -1),
            criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        )

        obj_points.append(objp)
        img_points.append(corners_subpix)

# 7. 标定相机（计算内参矩阵 K 和外参）
ret, K, dist, rvecs, tvecs = cv2.calibrateCamera(
    obj_points, img_points, gray.shape[::-1], None, None
)

print("内参矩阵 K:\n", K)

# 9. 点云展示及相机观察方向可视化
# 创建 3D 图形
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# 遍历每张图片对应的三维点和外参
for i in range(len(obj_points)):
    # 获取当前图片的旋转向量和平移向量
    rvec = rvecs[i]
    tvec = tvecs[i]

    # 将旋转向量转换为旋转矩阵
    R, _ = cv2.Rodrigues(rvec)

    # 提取当前图片的三维点
    obj_point = obj_points[i]

    # 提取 x, y, z 坐标
    x = obj_point[:, 0]
    y = obj_point[:, 1]
    z = obj_point[:, 2]

    # 绘制三维点
    ax.scatter(x, y, z, c='b', label='Feature Points' if i == 0 else "")

    # 相机在世界坐标系中的位置
    camera_position = tvec.flatten()

    # 相机的观察方向（假设相机朝向 Z 轴正方向）
    # camera_direction = (R @ np.array([0, 0, -1]))[0:1]
    camera_direction = []
    camera_direction = (R @ np.array([0, 0, 1]))[0:2]
    camera_direction = [*camera_direction,((R @ np.array([0, 0, -1]))[-1])]

    # 绘制相机位置
    ax.scatter(*camera_position, c='r', marker='s', label='Camera Position' if i == 0 else "")

    # 绘制相机观察方向
    ax.quiver(*camera_position, *camera_direction,length=40, color='g', label='Camera Direction' if i == 0 else "")

# ax.set_box_aspect([1, 1, 1])
ax.set_aspect('equal')

# 设置坐标轴标签
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

# 显示图例
ax.legend()

# 显示图形
plt.show()